#include "hybrid_rpc_server.hpp"
#include <iostream>
#include <cstring>
#include <chrono>

// 安全获取二进制ID（返回vector<uint8_t>）
std::vector<uint8_t> getBinaryIdentity (zmq_msg_t &identity_msg)
{
    uint8_t *data = static_cast<uint8_t *> (zmq_msg_data (&identity_msg));
    size_t size = zmq_msg_size (&identity_msg);
    return std::vector<uint8_t> (data, data + size);
}

// 二进制ID转可打印字符串（16进制）
std::string toHexString (const std::vector<uint8_t> &binary)
{
    static const char hex_digits[] = "0123456789ABCDEF";
    std::string result;
    result.reserve (binary.size () * 2);
    for (uint8_t byte : binary) {
        result.push_back (hex_digits[byte >> 4]);
        result.push_back (hex_digits[byte & 0xF]);
    }
    return result;
}

HybridRPCServer::HybridRPCServer (const std::string &rpc_bind_addr,
                                  const std::string &stream_bind_addr) :
    ctx_ (zmq_ctx_new ()), running_ (false)
{
    // 初始化ROUTER套接字
    rpc_socket_ = zmq_socket (ctx_, ZMQ_ROUTER);
    if (zmq_bind (rpc_socket_, rpc_bind_addr.c_str ()) != 0) {
        zmq_ctx_destroy (ctx_);
        throw std::runtime_error ("RPC bind failed: "
                                  + std::string (zmq_strerror (errno)));
    }

    // 初始化STREAM套接字
    stream_socket_ = zmq_socket (ctx_, ZMQ_STREAM);
    if (zmq_bind (stream_socket_, stream_bind_addr.c_str ()) != 0) {
        zmq_close (rpc_socket_);
        zmq_ctx_destroy (ctx_);
        throw std::runtime_error ("STREAM bind failed: "
                                  + std::string (zmq_strerror (errno)));
    }

    // 设置套接字超时（毫秒）
    int timeout = 100;
    zmq_setsockopt (rpc_socket_, ZMQ_RCVTIMEO, &timeout, sizeof (timeout));
    zmq_setsockopt (stream_socket_, ZMQ_RCVTIMEO, &timeout, sizeof (timeout));
}

HybridRPCServer::~HybridRPCServer ()
{
    stop ();
    zmq_close (rpc_socket_);
    zmq_close (stream_socket_);
    zmq_ctx_destroy (ctx_);
}

void HybridRPCServer::registerMethod (const std::string &name,
                                      std::function<json (const json &)> func)
{
    methods_[name] = func;
}

void HybridRPCServer::start ()
{
    std::lock_guard<std::mutex> lock (thread_mutex_);

    if (running_)
        return; // 防止重复启动

    running_ = true;
    worker_thread_ = std::make_unique<std::thread> ([this] () {
        std::cout << "Server worker thread started (ID: "
                  << std::this_thread::get_id () << ")\n";

        try {
            runEventLoop (); // 进入主事件循环
        }
        catch (const std::exception &e) {
            std::cerr << "Event loop crashed: " << e.what () << std::endl;
            running_ = false;
        }
    });

// 建议设置线程名（Linux/Mac）
#ifdef __linux__
    pthread_setname_np (worker_thread_->native_handle (), "zmq_server");
#endif
}

void HybridRPCServer::stop ()
{
    {
        std::lock_guard<std::mutex> lock (thread_mutex_);
        if (!running_)
            return;

        running_ = false; // 通知线程退出
    }

    if (worker_thread_ && worker_thread_->joinable ()) {
        worker_thread_->join (); // 等待线程结束
        worker_thread_.reset (); // 释放线程资源
        std::cout << "Server worker thread stopped\n";
    }
}

void HybridRPCServer::runEventLoop ()
{
    zmq_pollitem_t items[] = {{rpc_socket_, 0, ZMQ_POLLIN, 0},
                              {stream_socket_, 0, ZMQ_POLLIN, 0}};

    while (running_) {
        int rc = zmq_poll (items, 2, 100); // 100ms超时
        if (rc < 0) {
            std::cerr << "Poll error: " << zmq_strerror (errno) << std::endl;
            break;
        }

        if (items[0].revents & ZMQ_POLLIN) {
            handleRPCMessage ();
        }

        if (items[1].revents & ZMQ_POLLIN) {
            handleStreamMessage ();
        }
    }
}

void HybridRPCServer::handleRPCMessage ()
{
    zmq_msg_t identity, data;

    // 接收身份帧
    zmq_msg_init (&identity);
    if (zmq_msg_recv (&identity, rpc_socket_, 0) <= 0) {
        zmq_msg_close (&identity);
        return;
    }

    // 接收数据帧
    zmq_msg_init (&data);
    if (zmq_msg_recv (&data, rpc_socket_, 0) <= 0) {
        zmq_msg_close (&identity);
        zmq_msg_close (&data);
        return;
    }

    // 处理请求
    std::string request_str (static_cast<char *> (zmq_msg_data (&data)),
                             zmq_msg_size (&data));
    std::string response_str = processRPCRequest (request_str);
    zmq_msg_close (&data);

    // 发送响应
    zmq_msg_send (&identity, rpc_socket_, ZMQ_SNDMORE);
    zmq_msg_close (&identity);

    zmq_msg_init_size (&data, response_str.size ());
    memcpy (zmq_msg_data (&data), response_str.data (), response_str.size ());
    zmq_msg_send (&data, rpc_socket_, 0);
    zmq_msg_close (&data);
}

void HybridRPCServer::handleStreamMessage ()
{
    zmq_msg_t identity, data;

    // 1. 接收ID帧
    zmq_msg_init (&identity);
    if (zmq_msg_recv (&identity, stream_socket_, 0) <= 0) {
        zmq_msg_close (&identity);
        return;
    }

    // 2. 接收数据帧（非阻塞）
    zmq_msg_init (&data);
    int rc = zmq_msg_recv (&data, stream_socket_, ZMQ_DONTWAIT);

    // 3. 处理连接/断开事件
    std::string client_id (static_cast<char *> (zmq_msg_data (&identity)),
                           zmq_msg_size (&identity));
    // 安全获取二进制ID
    auto client_id_bin = getBinaryIdentity (identity);
    std::string client_id_hex = toHexString (client_id_bin);

    if (rc == -1 && errno == EAGAIN) {
        if (!client_id.empty ()) {
            std::lock_guard<std::mutex> lock (streams_mutex_);

            auto client_id_bin = getBinaryIdentity (identity);
            std::string client_id_hex = toHexString (client_id_bin);

            connected_streams_.erase (client_id_hex);
            std::cout << "STREAM Client disconnected: " << client_id_hex
                      << std::endl;
        }
        zmq_msg_close (&identity);
        return;
    }

    // 4. 处理数据或断开
    if (zmq_msg_size (&data) == 0) {
        // 断开连接
        std::lock_guard<std::mutex> lock (streams_mutex_);

        connected_streams_.insert (client_id_hex);
        std::cout << "STREAM Client connected: " << client_id_hex << std::endl;
    } else {
        // 处理数据
        std::string msg (static_cast<char *> (zmq_msg_data (&data)),
                         zmq_msg_size (&data));
        std::cout << "STREAM[" << client_id_hex << "] Received: " << msg
                  << std::endl;
        sendStreamResponse (client_id, msg);
    }

    zmq_msg_close (&identity);
    zmq_msg_close (&data);
}

std::string HybridRPCServer::processRPCRequest (const std::string &request)
{
    try {
        json request_json = json::parse (request);

        if (!request_json.contains ("method")
            || !request_json["method"].is_string ()) {
            return R"({" jsonrpc ":" 2.0 "," error ":{" code ":-32600," message ":" Invalid Request "}," id ":null})";
        }

        std::string method_name = request_json["method"];
        json params = request_json.value ("params", json::array ());

        auto it = methods_.find (method_name);
        if (it == methods_.end ()) {
            return R"({" jsonrpc ":" 2.0 "," error ":{" code ":-32601," message ":" Method not found "}," id
                                                               ":null})";
        }

        json result = it->second (params);
        json response = {{"jsonrpc", "2.0"},
                         {"result", result},
                         {"id", request_json.value ("id", json ())}};
        return response.dump ();
    }
    catch (const json::exception &e) {
        return R"({" jsonrpc ":" 2.0 "," error ":{" code ":-32700," message ":" Parse error "}," id ":null})";
    }
    catch (const std::exception &e) {
        return R"({" jsonrpc ":" 2.0 "," error ":{" code ":-32603," message ":" Internal error "}," id ":null})";
    }
}

void HybridRPCServer::sendStreamResponse (const std::string &client_id,
                                          const std::string &data)
{
    zmq_msg_t identity, msg;

    // 发送ID帧
    zmq_msg_init_size (&identity, client_id.size ());
    memcpy (zmq_msg_data (&identity), client_id.data (), client_id.size ());
    zmq_msg_send (&identity, stream_socket_, ZMQ_SNDMORE);
    zmq_msg_close (&identity);

    // 发送数据帧
    zmq_msg_init_size (&msg, data.size ());
    memcpy (zmq_msg_data (&msg), data.data (), data.size ());
    zmq_msg_send (&msg, stream_socket_, 0);
    zmq_msg_close (&msg);
}